import sys
import os
import json
import traceback
import warnings

# 设置 UTF-8 编码输出，避免中文乱码
import io
sys.stdout = io.TextIOWrapper(sys.stdout.buffer, encoding='utf-8')
sys.stderr = io.TextIOWrapper(sys.stderr.buffer, encoding='utf-8')

# 过滤警告信息，减少stderr输出
warnings.filterwarnings('ignore')
os.environ['PYTHONWARNINGS'] = 'ignore'

# 设置PaddlePaddle相关环境变量，减少日志输出和模型下载信息
os.environ['FLAGS_print_model_stats'] = '0'
os.environ['FLAGS_enable_parallel_graph'] = '0'
os.environ['GLOG_minloglevel'] = '2'  # 只显示错误级别日志
os.environ['GLOG_vmodule'] = 'inference_engine=0'

# 设置PaddleOCR相关环境变量，减少模型下载提示
os.environ['HUB_HOME'] = os.path.join(os.path.dirname(__file__), '.paddlehub')
os.environ['PADDLEHUB_HOME'] = os.path.join(os.path.dirname(__file__), '.paddlehub')

# 设置控制台代码页为UTF-8（Windows）
if os.name == 'nt':
    try:
        import locale
        locale.setlocale(locale.LC_ALL, 'C.UTF-8')
    except:
        try:
            # Windows fallback
            import codecs
            sys.stdout = codecs.getwriter('utf-8')(sys.stdout.detach())
            sys.stderr = codecs.getwriter('utf-8')(sys.stderr.detach())
        except:
            pass

try:
    import tkinter as tk
    from tkinter import messagebox
    HAS_TK = True
except ImportError:
    HAS_TK = False

def show_download_dialog(download_url):
    if HAS_TK:
        root = tk.Tk()
        root.withdraw()
        msg = f"未检测到本地独立Python环境！\n\n请下载绿色版Python及依赖压缩包，并解压到本插件目录下。\n\n下载地址：{download_url}\n\n解压后请重试。"
        messagebox.showerror("Python环境缺失", msg)
        root.destroy()
    else:
        print(json.dumps({'error': '未检测到本地独立Python环境，且tkinter不可用。请下载并解压到插件目录。'}))

PLUGIN_DIR = os.path.dirname(os.path.abspath(__file__))
PYTHON_DIR = os.path.join(PLUGIN_DIR, 'python')
PYTHON_EXE = os.path.join(PYTHON_DIR, 'python.exe')
MANIFEST_PATH = os.path.join(PLUGIN_DIR, 'manifest.json')

# 读取downurl
python_downurl = None
try:
    import json as _json
    with open(MANIFEST_PATH, 'r', encoding='utf-8') as f:
        manifest = _json.load(f)
        python_downurl = manifest.get('downurl')
except Exception as e:
    python_downurl = 'https://www.python.org/ftp/python/3.10.11/python-3.10.11-embed-amd64.zip'
    print(json.dumps({'warn': f'manifest读取失败: {str(e)}'}))

# 检查python.exe
if not os.path.exists(PYTHON_EXE):
    show_download_dialog(python_downurl)
    print(json.dumps({'error': '未检测到本地独立Python环境，请下载并解压到插件目录。'}))
    sys.exit(10)

# 只要python/Lib/site-packages存在即可
SITE_PACKAGES = os.path.join(PYTHON_DIR, 'Lib', 'site-packages')
if not os.path.exists(SITE_PACKAGES):
    print(json.dumps({'error': '未检测到site-packages目录，请用pip安装依赖。'}))
    sys.exit(11)

# 兼容绿色版/嵌入式python：强制将 site-packages 加入 sys.path
if SITE_PACKAGES not in sys.path:
    sys.path.insert(0, SITE_PACKAGES)

# urllib3 2.x 在部分依赖下会有类型注解兼容性问题，强制降级为 1.26.x，降级后自动退出，需重启脚本
import subprocess
try:
    import urllib3
    v = tuple(map(int, urllib3.__version__.split('.')[:2]))
    if v >= (2, 0):
        print(json.dumps({'info': f'检测到 urllib3 版本为 {urllib3.__version__}，即将降级...'}))
        subprocess.check_call([PYTHON_EXE, '-m', 'pip', 'install', 'urllib3<2'])
        print(json.dumps({'info': '已自动降级urllib3到1.26.x，请重新运行本脚本！'}))
        sys.exit(99)
except Exception as e:
    print(json.dumps({'warn': f'urllib3降级检测/处理异常: {str(e)}'}))

# 检查 paddleocr 是否可导入，不可导入则自动安装并退出
try:
    import paddleocr
except ImportError:
    print(json.dumps({'info': '检测到 paddleocr 未安装，正在自动安装...'}))
    try:
        subprocess.check_call([PYTHON_EXE, '-m', 'pip', 'install', 'paddleocr', 'paddlepaddle'])
        print(json.dumps({'info': '已自动安装 paddleocr 及依赖，请重新运行本脚本！'}))
    except Exception as e:
        print(json.dumps({'error': f'paddleocr 自动安装失败: {str(e)}，请手动运行 python.exe -m pip install paddleocr paddlepaddle'}))
    sys.exit(98)

# 调试：输出 sys.path，确认依赖查找路径
# print(json.dumps({'debug_sys_path': sys.path}))  # 注释掉调试输出

try:
    from paddleocr import PaddleOCR
except ImportError:
    print(json.dumps({'error': 'PaddleOCR未安装，请运行 python.exe -m pip install paddleocr paddlepaddle'}))
    sys.exit(1)

def safe_log_result(result):
    # 兼容 ndarray、bytes、常见不可序列化类型，便于调试
    try:
        import numpy as np
    except ImportError:
        np = None
    try:
        if isinstance(result, (list, dict, str, int, float, bool)):
            return json.dumps({'raw_result': result}, ensure_ascii=False)
        if np and isinstance(result, np.ndarray):
            return json.dumps({'raw_result': result.tolist(), 'shape': result.shape, 'dtype': str(result.dtype)}, ensure_ascii=False)
        if isinstance(result, bytes):
            return json.dumps({'raw_result': result.decode(errors='replace')}, ensure_ascii=False)
        return json.dumps({'raw_result': str(result)}, ensure_ascii=False)
    except Exception as e:
        return json.dumps({'warn': f'raw_result无法序列化: {str(e)}'})

def parse_ocr_result(result):
    """解析 PaddleOCR 的结果，支持多种格式"""
    try:
        import numpy as np
    except ImportError:
        np = None
        
    text_lines = []
    
    # print(json.dumps({'debug': f'解析 OCR 结果，类型: {type(result)}'}, ensure_ascii=False))  # 注释掉调试输出
    
    try:
        # 1. 兼容新版 predict API，支持 paddlex 文档结构化输出（list[dict]）
        if isinstance(result, list) and result and isinstance(result[0], dict):
            # print(json.dumps({'debug': '使用 list[dict] 格式解析'}, ensure_ascii=False))  # 注释掉调试输出
            for item in result:
                if 'rec_texts' in item and isinstance(item['rec_texts'], list):
                    for txt in item['rec_texts']:
                        if isinstance(txt, str) and len(txt.strip()) > 0:
                            text_lines.append(txt.strip())
                            
        # 2. 兼容新版 predict API（dict）
        elif isinstance(result, dict):
            # print(json.dumps({'debug': '使用 dict 格式解析'}, ensure_ascii=False))  # 注释掉调试输出
            if 'rec_texts' in result and isinstance(result['rec_texts'], list):
                for txt in result['rec_texts']:
                    if isinstance(txt, str) and len(txt.strip()) > 0:
                        text_lines.append(txt.strip())
            elif 'text' in result:
                if isinstance(result['text'], str):
                    text_lines.append(result['text'].strip())
                elif isinstance(result['text'], list):
                    for txt in result['text']:
                        if isinstance(txt, str) and len(txt.strip()) > 0:
                            text_lines.append(txt.strip())
                            
        # 3. 兼容传统 OCR 结果格式 - 三层嵌套列表
        elif isinstance(result, list):
            # print(json.dumps({'debug': f'使用 list 格式解析，长度: {len(result)}'}, ensure_ascii=False))  # 注释掉调试输出
            
            # 处理 [[[x1,y1],[x2,y2],...], (text, confidence)] 格式
            for block in result:
                if isinstance(block, (list, tuple)) and len(block) >= 2:
                    # block[0] 是坐标，block[1] 是文本信息
                    text_info = block[1]
                    if isinstance(text_info, (list, tuple)) and len(text_info) >= 1:
                        text_content = text_info[0]
                        if isinstance(text_content, str) and len(text_content.strip()) > 0:
                            text_lines.append(text_content.strip())
                    elif isinstance(text_info, str) and len(text_info.strip()) > 0:
                        text_lines.append(text_info.strip())
                        
        # 4. 兼容 numpy array 结果
        elif np and isinstance(result, np.ndarray):
            # print(json.dumps({'debug': '使用 numpy array 格式解析'}, ensure_ascii=False))  # 注释掉调试输出
            try:
                arr = result.tolist()
                if isinstance(arr, list):
                    for item in arr:
                        if isinstance(item, str) and len(item.strip()) > 0:
                            text_lines.append(item.strip())
            except Exception as e:
                # print(json.dumps({'warn': f'numpy array 解析异常: {str(e)}'}, ensure_ascii=False))  # 注释掉调试输出
                pass
                
        # 5. 如果还是没有解析到文本，尝试直接转换
        if not text_lines:
            # print(json.dumps({'debug': '使用兜底字符串转换'}, ensure_ascii=False))  # 注释掉调试输出
            result_str = str(result)
            if len(result_str) > 10:  # 避免返回无意义的短字符串
                text_lines.append(result_str)
                
        # 6. 清理和去重
        cleaned_lines = []
        for line in text_lines:
            if line and len(line.strip()) > 0:
                # 过滤掉明显不是文本的内容
                cleaned_line = line.strip()
                if any(c.isalnum() or '\u4e00' <= c <= '\u9fff' for c in cleaned_line):
                    cleaned_lines.append(cleaned_line)
                    
        # 7. 去重并合并
        unique_lines = list(dict.fromkeys(cleaned_lines))  # 保持顺序去重
        
        final_text = '\n'.join(unique_lines) if unique_lines else ''
        
        # print(json.dumps({
        #     'debug': f'解析完成，原始行数: {len(text_lines)}, 清理后: {len(cleaned_lines)}, 去重后: {len(unique_lines)}'
        # }, ensure_ascii=False))  # 注释掉调试输出
        
        return final_text
        
    except Exception as e:
        # print(json.dumps({'error': f'OCR 结果解析异常: {str(e)}'}, ensure_ascii=False))  # 注释掉调试输出
        return str(result) if result else ''

def main():
    try:
        # 输出启动信息
        print(json.dumps({'info': '开始OCR处理...'}, ensure_ascii=False), flush=True)
        
        if len(sys.argv) < 3:
            print(json.dumps({'error': '参数不足，需传入图片路径和语言（ch/en）'}, ensure_ascii=False))
            sys.exit(2)
            
        image_path = sys.argv[1]
        lang = sys.argv[2]
        
        print(json.dumps({'info': f'图片路径={image_path}, 语言={lang}'}, ensure_ascii=False), flush=True)
        
        # 验证图片文件存在
        if not os.path.exists(image_path):
            print(json.dumps({'error': f'图片文件不存在: {image_path}'}, ensure_ascii=False))
            sys.exit(4)
            
        # print(json.dumps({'info': '开始初始化 PaddleOCR...'}, ensure_ascii=False))  # 注释掉调试输出
        
        # 初始化 PaddleOCR - 使用渐进式参数添加策略
        base_kwargs = {
            'use_angle_cls': True,
            'lang': lang,
        }
        
        # 按优先级尝试不同的参数组合
        param_combinations = [
            # 完整参数组合（最新版本）
            {**base_kwargs, 'show_log': False, 'use_textline_orientation': True},
            # 去掉 use_textline_orientation
            {**base_kwargs, 'show_log': False},
            # 去掉 show_log 
            {**base_kwargs, 'use_textline_orientation': True},
            # 最基础参数（最兼容）
            base_kwargs
        ]
        
        ocr = None
        init_error = None
        used_params = None
        
        # 在初始化前先输出状态信息
        print(json.dumps({'info': 'PaddleOCR 初始化中，请稍候...'}, ensure_ascii=False), flush=True)
        
        for i, kwargs in enumerate(param_combinations):
            try:
                # 临时禁用stderr输出，避免模型下载信息干扰
                import contextlib
                with open(os.devnull, 'w') as devnull:
                    with contextlib.redirect_stderr(devnull):
                        ocr = PaddleOCR(**kwargs)
                
                used_params = kwargs
                print(json.dumps({'info': f'PaddleOCR 初始化成功，使用参数组合 {i+1}: {list(kwargs.keys())}'}, ensure_ascii=False), flush=True)
                break
            except (TypeError, ValueError) as e:
                init_error = str(e)
                if 'show_log' in str(e) or 'use_textline_orientation' in str(e):
                    # 参数不支持，尝试下一个组合
                    continue
                else:
                    # 其他参数错误，也尝试下一个组合
                    continue
            except Exception as e:
                init_error = str(e)
                # 对于其他异常，记录但继续尝试
                print(json.dumps({'warn': f'参数组合 {i+1} 初始化失败: {str(e)}'}, ensure_ascii=False), flush=True)
                continue
        
        if ocr is None:
            print(json.dumps({'error': f'PaddleOCR初始化失败，所有参数组合都不支持。最后错误: {init_error}'}, ensure_ascii=False))
            sys.exit(5)
        
        # 可选：输出PaddleOCR版本信息用于调试
        try:
            import paddlepaddle as paddle
            paddle_version = paddle.__version__
            # print(json.dumps({'info': f'使用 PaddlePaddle 版本: {paddle_version}'}, ensure_ascii=False))
        except:
            # print(json.dumps({'info': 'PaddlePaddle 版本信息获取失败'}, ensure_ascii=False))
            pass
                
        # print(json.dumps({'info': 'PaddleOCR 初始化完成'}, ensure_ascii=False))  # 注释掉调试输出
        
        # 执行 OCR 识别
        print(json.dumps({'info': '开始 OCR 识别...'}, ensure_ascii=False), flush=True)
        
        result = None
        text = ''
        
        try:
            # 优先尝试使用 ocr 方法（更稳定）
            if hasattr(ocr, 'ocr'):
                print(json.dumps({'info': '使用 ocr 方法进行识别'}, ensure_ascii=False), flush=True)
                
                # 添加超时和异常处理
                import signal
                import time
                
                def timeout_handler(signum, frame):
                    raise TimeoutError("OCR识别超时")
                
                # 设置超时（60秒）
                if os.name != 'nt':  # 非Windows系统
                    signal.signal(signal.SIGALRM, timeout_handler)
                    signal.alarm(60)
                
                try:
                    start_time = time.time()
                    result = ocr.ocr(image_path)
                    end_time = time.time()
                    print(json.dumps({'info': f'OCR识别耗时: {end_time - start_time:.2f}秒'}, ensure_ascii=False), flush=True)
                    
                    if os.name != 'nt':
                        signal.alarm(0)  # 取消超时
                        
                except TimeoutError:
                    print(json.dumps({'error': 'OCR识别超时（60秒），请检查图片大小或网络连接'}, ensure_ascii=False))
                    sys.exit(7)
                except Exception as ocr_error:
                    print(json.dumps({'warn': f'ocr方法执行异常: {str(ocr_error)}'}, ensure_ascii=False), flush=True)
                    result = None  # 标记为失败，将尝试backup方法
                
                if result is not None:
                    print(json.dumps({'debug': f'ocr 方法执行完成，结果类型: {type(result)}'}, ensure_ascii=False), flush=True)
                    text = parse_ocr_result(result)
                    print(json.dumps({'info': f'文本解析完成，长度: {len(text) if text else 0}'}, ensure_ascii=False), flush=True)
                    
            elif hasattr(ocr, 'predict'):
                print(json.dumps({'info': '使用 predict 方法进行识别'}, ensure_ascii=False), flush=True)
                result = ocr.predict(image_path)
                print(json.dumps({'debug': f'predict 方法执行完成，结果类型: {type(result)}'}, ensure_ascii=False), flush=True)
                text = parse_ocr_result(result)
            else:
                raise AttributeError('PaddleOCR对象没有ocr或predict方法')
                
        except Exception as e:
            error_msg = str(e)
            print(json.dumps({'warn': f'第一种方法失败: {error_msg}'}, ensure_ascii=False), flush=True)
            
            # 如果是第一种方法失败，尝试备用方法
            if hasattr(ocr, 'predict') and ('ocr' in error_msg.lower() or result is None):
                try:
                    print(json.dumps({'info': '尝试使用predict方法作为备用'}, ensure_ascii=False), flush=True)
                    result = ocr.predict(image_path)
                    text = parse_ocr_result(result)
                except Exception as e2:
                    print(json.dumps({'error': f'OCR执行异常 (两种方法都失败): ocr方法={error_msg}, predict方法={str(e2)}'}, ensure_ascii=False))
                    sys.exit(6)
            elif hasattr(ocr, 'ocr') and 'predict' in error_msg.lower():
                try:
                    print(json.dumps({'info': '尝试使用ocr方法作为备用'}, ensure_ascii=False), flush=True)
                    result = ocr.ocr(image_path)
                    text = parse_ocr_result(result)
                except Exception as e2:
                    print(json.dumps({'error': f'OCR执行异常 (两种方法都失败): predict方法={error_msg}, ocr方法={str(e2)}'}, ensure_ascii=False))
                    sys.exit(6)
            else:
                print(json.dumps({'error': f'OCR执行异常: {error_msg}'}, ensure_ascii=False))
                sys.exit(6)
        
        # 确保有文本输出
        if not text or text.strip() == '':
            text = '未识别到文本内容'
            
        # print(json.dumps({'info': f'OCR 识别完成，文本长度: {len(text)}'}, ensure_ascii=False))  # 注释掉调试输出
        
        # 输出最终结果
        final_result = {
            'text': text,
            'success': True,
            'image_path': image_path,
            'lang': lang
        }
        
        print(json.dumps(final_result, ensure_ascii=False))
        
    except Exception as e:
        error_msg = f'OCR 处理异常: {str(e)}'
        print(json.dumps({
            'error': error_msg,
            'trace': traceback.format_exc(),
            'success': False
        }, ensure_ascii=False))
        sys.exit(3)

if __name__ == '__main__':
    main()